from langchain.agents import create_agent
from langchain_deepseek import ChatDeepSeek
from langchain_core.messages import HumanMessage
import json


def book_hotel(hotel_name: str):
    """预订酒店"""
    return f"成功预订入住 {hotel_name}."


def book_flight(from_airport: str, to_airport: str):
    """预订航班"""
    return f"成功预订了从 {from_airport} 到 {to_airport}."


# 初始化DeepSeek模型客户端，并配置相关参数以控制生成行为
llm = ChatDeepSeek(
    model="deepseek-chat",  # 指定使用的模型名称
    api_key="sk-ddad02a64c4c47ea9f6f526ef47cb602",  # 替换为你的API Key
    temperature=0.7,  # 控制生成结果的随机性，0表示更确定性的输出
    max_tokens=None,  # 生成内容的最大长度，None代表使用模型默认值
    timeout=60,  # 请求超时设置（单位秒）
    max_retries=2,  # 请求失败时最大重试次数
)

# 创建专门的代理 - 使用正确的参数名 system_prompt
flight_assistant = create_agent(
    model=llm,
    tools=[book_flight],
    system_prompt="您是一名航班预订助理。只负责处理航班预订相关任务。当用户提到航班、飞机、机场等相关内容时，使用 book_flight 工具。",
    name="flight_assistant"
)

hotel_assistant = create_agent(
    model=llm,
    tools=[book_hotel],
    system_prompt="您是酒店预订助理。只负责处理酒店预订相关任务。当用户提到酒店、入住、住宿等相关内容时，使用 book_hotel 工具。",
    name="hotel_assistant"
)


# 创建协调器（替代 langgraph_supervisor）
class SimpleCoordinator:
    def __init__(self, flight_agent, hotel_agent, llm):
        self.flight_agent = flight_agent
        self.hotel_agent = hotel_agent
        self.llm = llm

    def analyze_request(self, user_input):
        """分析用户请求，确定需要哪些代理处理"""
        analysis_prompt = f"""
        分析以下用户请求，判断是否需要航班预订、酒店预订，或两者都需要：

        用户请求："{user_input}"

        请回答以下问题：
        1. 是否需要预订航班？（是/否）
        2. 是否需要预订酒店？（是/否）
        3. 航班相关信息（如出发地、目的地）
        4. 酒店相关信息（如酒店名称）

        请用JSON格式回答：
        {{
            "need_flight": true/false,
            "need_hotel": true/false,
            "flight_info": {{"from": "", "to": ""}},
            "hotel_info": {{"name": ""}}
        }}
        """

        try:
            response = self.llm.invoke(analysis_prompt)
            return json.loads(response.content)
        except:
            # 简单的关键词分析作为备选
            need_flight = any(word in user_input for word in ["航班", "飞机", "机场", "飞往", "从"])
            need_hotel = any(word in user_input for word in ["酒店", "入住", "住宿", "住"])

            return {
                "need_flight": need_flight,
                "need_hotel": need_hotel,
                "flight_info": {},
                "hotel_info": {}
            }

    def process_request(self, user_input):
        """处理用户请求"""
        print(f"用户请求: {user_input}")

        # 分析请求
        analysis = self.analyze_request(user_input)
        print(f"请求分析: {json.dumps(analysis, ensure_ascii=False, indent=2)}")

        results = {}

        # 处理航班预订
        if analysis["need_flight"]:
            print("\n--- 航班预订助理处理中 ---")
            try:
                # 调用航班代理处理用户输入，获取航班预订结果
                # 将用户输入封装为HumanMessage对象并传递给flight_agent进行处理
                flight_result = self.flight_agent.invoke({
                    "messages": [HumanMessage(content=user_input)]
                })

                # 从航班代理的响应中提取最终的回复内容
                # 如果返回结果是字典类型，则提取messages列表中最后一条消息的内容
                # 否则将整个结果转换为字符串形式
                results["flight"] = flight_result.get("messages", [])[-1].content if isinstance(flight_result,
                                                                                                dict) else str(
                    flight_result)

                # 打印航班预订的最终结果
                print(f"航班预订结果: {results['flight']}")

            except Exception as e:
                print(f"航班预订出错: {e}")
                results["flight"] = f"航班预订出错: {e}"

        # 处理酒店预订
        if analysis["need_hotel"]:
            print("\n--- 酒店预订助理处理中 ---")
            try:
                hotel_result = self.hotel_agent.invoke({
                    "messages": [HumanMessage(content=user_input)]
                })
                results["hotel"] = hotel_result.get("messages", [])[-1].content if isinstance(hotel_result,
                                                                                              dict) else str(
                    hotel_result)
                print(f"酒店预订结果: {results['hotel']}")
            except Exception as e:
                print(f"酒店预订出错: {e}")
                results["hotel"] = f"酒店预订出错: {e}"

        # 综合结果
        if len(results) > 0:
            print("\n=== 最终处理结果 ===")
            for key, value in results.items():
                print(f"{key}: {value}")
        else:
            print("无法识别具体的预订需求，请明确说明您需要预订航班还是酒店。")

        return results


# 创建协调器实例
coordinator = SimpleCoordinator(flight_assistant, hotel_assistant, llm)

# 处理用户请求
if __name__ == "__main__":
    user_request = "预订从波士顿飞往肯尼迪机场的航班，并入住麦基特里克酒店"
    coordinator.process_request(user_request)
